# a simple model of epidemic spread with spatial segregation; generic model template used with thanks to Rick Riolo
# Jon Zelner (jzelner at umich dot edu)
# 8/2009
# University of Michigan

try:
    import scipy as sp
    import numpy as np
    import pylab as pl
except:
    raise("You must have scipy, numpy and pylab installed to run this model!")

import heapq
import copy

params = {'beta':10.0, 'gamma':1.0, 'nX': 10, 'nY':10, 'nSize':1000, 'ieSize':10, 'sc': .05, 'tau1': 0.0, 'tau2': 0.0, 'rho': 3.0, 'gw':-2.0, 'gb':-2.0 }


class sirNeighborhood():
#neighborhoods are in our most basic unit; the closest thing this model has to an agent.
#Individuals mix homogeneously within neighborhoods. Users can subclass this most basic neighborhood to 
#incorporate different pathogen natural histories or within-neighborhood mixing dynamics.
    
    def __init__(self, model, S = 100, I = 0, R = 0, x = 0, y = 0, b = .5):
        self.model = model
        
        #initial conditions
        self.S = S #number of susceptible individuals
        self.I = I #number of infectious individuals
        self.R = R #number recovered
        self.N = S + I + R #total population size
        
        self.x = x #these let the neighborhood know its position on the grid 
                   #for calculating geographic distance from all other neighborhoods
        self.y = y
        
        self.b = b #pathogen transmissibility in E(# new cases)/infected/timestep
        
        #keep track of agent recovery times here; we'll use a heap to keep these in order
        self.recoveryTimeList = []
        
        #let the neighborhood know its own ID
        self.ID = self.model.addNeighborhood(self)
        
        #let the model know about the 
        self.model.FOIList[self.ID] = float(self.I) / float(self.N)
        
        
    
    def infect(self, numToI = 1):
        
        #can infect up to the whole neighborhood on a single timestep;
        #decrement number of susceptibles, increase # of infecteds
        self.S -= numToI
        self.I += numToI
        
        #tell the model abotu this change as well
        self.model.S -= numToI
        self.model.I += numToI
        
        #update the force of infection
        self.model.FOIList[self.ID] = float(self.I) / float(self.N)
        
        #we draw recovery times instead of drawing a random float on each step; this saves processor cycles
        #and also allows us to use a wider range of recovery time distributions.
        for i in xrange(numToI):
            self.drawRecoveryTime()
    
    
    
    
    def recover(self, numToR = 1):
        
        self.I -= numToR
        self.R += numToR
        
        self.model.I -= numToR
        self.model.R += numToR
        
        self.model.FOIList[self.ID] = float(self.I) / float(self.N)

    
    def drawRecoveryTime(self):
        #pushes a new recovery time onto the neighborhood's recovery heap
        heapq.heappush(self.recoveryTimeList, self.model.t + sp.random.exponential(self.model.gamma))
    
    def step(self):
        #find out how many individuals to infect on this step        
        numToI = self.transmissionKernel()
        
        numToR = 0
        #check to see if anyone is scheduled to recover this step
        while( len(self.recoveryTimeList) > 0  and self.recoveryTimeList[0] <= self.model.t):
            numToR += 1
            heapq.heappop(self.recoveryTimeList)
        
        #if there are individuals who need to recover this step, change their state to 'recovered'
        self.recover(numToR = numToR)
        
        #ditto for infecting individuals
        self.infect(numToI = numToI)
        
        
    
    def transmissionKernel(self):
        #FOI = Force Of Infection
        internalFOI = (1-self.model.sc) * self.model.beta * (float(self.I) / float(self.N))
        
        externalFOI = 0
       #print('FOIStart', externalFOI)
        for extID, extFOI in enumerate(self.model.FOIList):
            if self.ID != extID:
                #have an option to either incorporate social coupling or just look at mixing as a function of distance; this implementation is kind of 
                #kludgy and is basically begging to be improved
                if not self.model.socialModel :
                    externalFOI += extFOI * sp.math.exp(-self.model.spatialCouplingMatrix[self.ID, extID])
                else:
                    externalFOI += extFOI * self.model.spatialCouplingMatrix[self.ID, extID]
        
        #take a look and see what's coming from outside
        externalFOI = self.model.beta * self.model.sc  * externalFOI
        FOI = internalFOI + externalFOI
        
        #print('FOIEND', FOI)
            
        randDraws = sp.random.uniform(size = self.S)
        numToI = len(pl.find(randDraws <= FOI))
        
        return numToI
        
        

class model():
    
    def __init__(self, params, neighborhoodType = sirNeighborhood):
        
        self.params = params
        
        #by pasing a specific type of agent, we can try out different pathogen natural histories
        #without changing the mechanical structure of the model
        self.neighborhood = neighborhoodType
        
        self.beta = self.params['beta'] / 24.0
        self.gamma = self.params['gamma'] * 24.0
        self.xGridDim = self.params['nX']
        self.yGridDim = self.params['nY']
        self.neighborhoodSize = self.params['nSize']
        self.initEventSize = self.params['ieSize']
        self.totalNeighborhoods = self.xGridDim * self.yGridDim
        
        self.sc = self.params['sc']
        self.tau1 = self.params['tau1']
        self.tau2 = self.params['tau2']
        self.rho = self.params['rho']
        self.gw = self.params['gw']
        self.gb = self.params['gb']
        
        self.socialModel = 0
        
        self.spatialCouplingMatrix = np.ndarray([self.totalNeighborhoods, self.totalNeighborhoods])
        
        #here's where we keep track of  neighborhood instances
        self.neighborhoodList = []
        
        #here, we keep track of the states of the whole population of agents
        self.S = 0
        self.I = 0
        self.R = 0
        
        self.neighborhoodID = -1
        
        self.t = 0
        
        self.FOIList = np.empty(self.xGridDim*self.yGridDim)
    
    def setupSocialModel(self):
        self.socialModel = 1
        #create neighborhoods:
        i = -1
        for x in xrange(self.xGridDim):
            for y in xrange(self.yGridDim):
                i += 1
                if ( x < 3) and (y > 3 and y < 7):

                    self.neighborhood(self, S = self.neighborhoodSize, x = x, y = y, b = 1.0)
                    nI = self.neighborhoodList[i]

                else:      
                    self.neighborhood(self, S = self.neighborhoodSize, x = x, y = y, b = 0.0)
        
        nI.infect(numToI = self.initEventSize)
        #neighborhoodToInfect = sp.random.randint( len( self.neighborhoodList ) )      
        #nI = self.neighborhoodList[neighborhoodToInfect]
        self.initXY = (nI.y+.5, nI.x+.5)
        self.setupSocialSpatialCoupling()
    
    def setup(self):
        
        #create neighborhoods:
        for x in xrange(self.xGridDim):
            for y in xrange(self.yGridDim):
                self.neighborhood(self, S = self.neighborhoodSize, x = x, y = y)
        
        neighborhoodToInfect = sp.random.randint( len( self.neighborhoodList ) )      
        self.neighborhoodList[neighborhoodToInfect].infect(  numToI = self.initEventSize )
        nI = self.neighborhoodList[neighborhoodToInfect]
        self.initXY = (nI.x+.5, nI.y+.5)
        self.setupSpatialCoupling()
    
    
    def setupSpatialCoupling(self):
        
        #iterate over all neighborhoods and insert the euclidean distance between them into spatial coupling matrix (for now);
        #going to expand this so that it covers social covariates later
        
        for i in xrange(len(self.neighborhoodList)):
            for j in xrange(i, len(self.neighborhoodList)):
                #print(i,j)
                if ( i == j ):
                    self.spatialCouplingMatrix[i,j] = 0.0
                else:
                    nA = self.neighborhoodList[i]
                    nB = self.neighborhoodList[j]
                    
            
                    dist = sp.math.sqrt((nA.x - nB.x)**2 + (nA.y - nB.y)**2)
                    
                    self.spatialCouplingMatrix[i, j] = dist
                    self.spatialCouplingMatrix[j, i] = dist
            
    def setupSocialSpatialCoupling(self):
           
        #iterate over all neighborhoods and insert the euclidean distance between them into spatial coupling matrix (for now);
        #going to expand this so that it covers social covariates later
        
        for i in xrange(len(self.neighborhoodList)):
            neighborhoodA = self.neighborhoodList[i]
            for j in xrange(i, len(self.neighborhoodList)):
                #print(i,j) 
                if ( i == j ):
                    self.spatialCouplingMatrix[i,j] = 0.0
                else:
                    neighborhoodB = self.neighborhoodList[j]
                    
                    bA = neighborhoodA.b
                    bB = neighborhoodB.b
                    
                    nA = neighborhoodA.N
                    nB = neighborhoodB.N
                    
                    
                    qAB = (1-bA)*(sp.math.exp(self.gw*bB)) + bA*(sp.math.exp(self.gb*(1-bB)))
                    qBA = (1-bB)*(sp.math.exp(self.gw*bA)) + bB*(sp.math.exp(self.gb*(1-bA)))
                    

                    
                    euclidDist = sp.math.sqrt((neighborhoodA.x - neighborhoodB.x)**2 + (neighborhoodA.y - neighborhoodB.y)**2)
                    
                    totalDistAB = ((nA**self.tau1)*(nB**self.tau2)*qAB) * sp.math.exp(-euclidDist*self.rho)
                    totalDistBA = ((nA**self.tau1)*(nB**self.tau2)*qBA) * sp.math.exp(-euclidDist*self.rho)
                    

                    self.spatialCouplingMatrix[i, j] = totalDistAB
                    self.spatialCouplingMatrix[j, i] = totalDistBA
    
    def isOutbreak(self):
        N = self.S + self.I + self.R
        outbreak = 0
        while (self.I > 0):
            self.step()
            self.t += 1
            if ((self.I / float(N)) >= .05):
                outbreak = 1
                break
        return outbreak

    
    def runUntilDieOut(self, printOut = True):
        
        numI = []
        numS = []
        numR = []
        iGridList = []
        
        #as long as anyone is still infectious, keep the model running
        while(self.I > 0):
            #step all agents
            self.step()
            self.t += 1
            #bookkeeping
            numI.append(self.I)
            numS.append(self.S)
            numR.append(self.R)
            
            if printOut:
                print('t= ', self.t,'S', self.S, 'I', self.I, 'R', self.R)
            
            if (self.t % 96 == 0):
                iGridTemp = np.ndarray([self.xGridDim, self.yGridDim])
                for n in self.neighborhoodList:
                    iGridTemp[n.x, n.y] = float(n.I) / n.N
                    
                iGridList.append(iGridTemp)
        
        #now make grid for heatmap of level of infection in each neighborhood
        returnGrid = np.ndarray([self.xGridDim, self.yGridDim])
        compGrid = np.ndarray([self.xGridDim, self.yGridDim])
        
        
        for n in self.neighborhoodList:
            returnGrid[n.x, n.y] = float(n.R) / n.N
            compGrid[n.x, n.y] = n.b
        
        iGridList.append(returnGrid)
        
        return [numI, numS, numR, returnGrid, compGrid, iGridList]
    
    def addNeighborhood(self, neighborhood):
        
        self.neighborhoodList.append(neighborhood)
        
        self.S += neighborhood.S
        self.I += neighborhood.I
        self.R += neighborhood.R
        
        self.neighborhoodID += 1
        
        return self.neighborhoodID

    def step(self):
        for neighborhood in self.neighborhoodList:
            neighborhood.step()



def removeFromArray(array, value):
    array = np.delete(array, mlab.find(array = value))
    return array

def panelPcolorPlot(inputList):
    numPanels = len(inputList)
    pl.figure()
    nCol = 3.0
    nRow = sp.ceil(numPanels/nCol)
    
    pl.subplot(nRow,nCol,1)
    p = 0
    for i in xrange(numPanels):
        p += 1
        pl.subplot(nRow,nCol,p)
        pl.pcolor(inputList[i], vmin = 0, vmax = 1)
    
    pl.show()
        
    
    
        